{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os \n",
    "import numpy as np \n",
    "import json \n",
    "import sys \n",
    "import tqdm\n",
    "import trimesh \n",
    "sys.path.append(\"..\")\n",
    "from src.utils import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# collect all ycb models that has .sdf files\n",
    "ycb_dir = '/media/junting/SSD_data/ycb'\n",
    "model_ids = os.listdir(ycb_dir)\n",
    "metadata = {\n",
    "    'objects': {},\n",
    "}\n",
    "output_dir = '../../data/ycb'\n",
    "if not os.path.exists(output_dir):\n",
    "    os.makedirs(output_dir)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "for model_id in model_ids:\n",
    "    model_name = model_id[4:]\n",
    "    model_folder = os.path.join(ycb_dir, model_id)\n",
    "    sdf_file = os.path.join(model_folder, f'{model_name}.sdf')\n",
    "    \n",
    "    if os.path.exists(sdf_file):\n",
    "        metadata['objects'][model_id] = {\n",
    "            'model_id': model_id,\n",
    "            'model_name': model_name\n",
    "        }\n",
    "\n",
    "        # calculate the dimension of the object mesh\n",
    "        # Check if there are Google meshes; else use the TSDF folder\n",
    "        if \"google_16k\" in os.listdir(model_folder):\n",
    "            mesh_type = \"google_16k\"\n",
    "        else:\n",
    "            mesh_type = \"tsdf\"\n",
    "\n",
    "        if mesh_type == \"google_16k\":\n",
    "            mesh_file = os.path.join(model_folder, \"google_16k\", \"textured.obj\")\n",
    "        elif mesh_type == \"tsdf\":\n",
    "            mesh_file = os.path.join(model_folder, \"tsdf\", \"textured.obj\")\n",
    "        mesh = trimesh.load(mesh_file)\n",
    "        dimension = getMeshDimensions(mesh)\n",
    "        center = getMeshCenter(mesh)\n",
    "        metadata['objects'][model_id]['bbox_size'] = dimension.tolist()\n",
    "        metadata['objects'][model_id]['bbox_center'] = center.tolist()\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# filter out objects that are too small, i.e. < 4cm\n",
    "for model_id in list(metadata['objects'].keys()):\n",
    "    bbox_size = metadata['objects'][model_id]['bbox_size']\n",
    "    if np.min(bbox_size) < 0.05:\n",
    "        del metadata['objects'][model_id]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# manually remove skillet since too difficult to grasp for a gripper\n",
    "del metadata['objects']['027_skillet']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of objects: 27\n"
     ]
    }
   ],
   "source": [
    "print(f\"Number of objects: {len(metadata['objects'])}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# save metadata\n",
    "# with open(os.path.join(output_dir, 'metadata.json'), 'w') as f:\n",
    "#     json.dump(metadata, f, indent=4)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate thumbnail for each model "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "0it [00:00, ?it/s]\n"
     ]
    }
   ],
   "source": [
    "# Use external tool f3d to render thumbnails\n",
    "for model_id in tqdm.tqdm(metadata['objects'].keys()):\n",
    "    model_name = metadata['objects'][model_id]['model_name']\n",
    "    model_folder = os.path.join(ycb_dir, model_id)\n",
    "    if \"google_16k\" in os.listdir(model_folder):\n",
    "        mesh_file = os.path.join(model_folder, \"google_16k\", \"textured.obj\")\n",
    "        material_file = os.path.join(model_folder, \"google_16k\", \"textured_map.png\")\n",
    "    elif \"tsdf\" in os.listdir(model_folder):\n",
    "        mesh_file = os.path.join(model_folder, \"tsdf\", \"textured.obj\")\n",
    "        material_file = os.path.join(model_folder, \"tsdf\", \"textured.png\")\n",
    "    \n",
    "    output_file = os.path.join(output_dir, \"object_thumbnails\", f\"{model_name}.png\")\n",
    "\n",
    "    cmd = f\"/usr/bin/f3d --input={mesh_file} --output={output_file} --texture-base-color={material_file}\"\n",
    "    os.system(cmd)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Make collage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# collect all thumbnails in container_thumbnails folder, and collage as one image\n",
    "from PIL import Image\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.image as mpimg\n",
    "import numpy as np\n",
    "import os\n",
    "import math\n",
    "\n",
    "def collage(thumbnail_dir, output_path):\n",
    "\n",
    "    # Get a list of all the thumbnail file names in the folder\n",
    "    thumbnail_files = os.listdir(thumbnail_dir)\n",
    "\n",
    "    # Calculate the number of thumbnails in the folder\n",
    "    num_thumbnails = len(thumbnail_files)\n",
    "\n",
    "    # Calculate the size of the square image needed to fit all the thumbnails\n",
    "    image_size = int(math.ceil(math.sqrt(num_thumbnails)))\n",
    "\n",
    "    # Create a new blank image to use as the canvas for the collage\n",
    "    canvas = Image.new('RGB', (image_size * 100, image_size * 100), (255, 255, 255))\n",
    "\n",
    "    # Loop through each thumbnail file and paste it onto the canvas\n",
    "    for i in range(num_thumbnails):\n",
    "        # Open the thumbnail image\n",
    "        thumbnail = Image.open(os.path.join(thumbnail_dir, thumbnail_files[i]))\n",
    "        \n",
    "        # Resize the thumbnail to 100x100 pixels\n",
    "        thumbnail = thumbnail.resize((100, 100))\n",
    "        \n",
    "        # Calculate the x and y coordinates of where to paste the thumbnail onto the canvas\n",
    "        x = (i % image_size) * 100\n",
    "        y = (i // image_size) * 100\n",
    "        \n",
    "        # Paste the thumbnail onto the canvas\n",
    "        canvas.paste(thumbnail, (x, y))\n",
    "\n",
    "    # Save the canvas image\n",
    "    canvas.save(output_path)\n",
    "\n",
    "# save object thumbnails as one image\n",
    "thumbnail_dir = os.path.join(output_dir, \"object_thumbnails\")\n",
    "object_collage_path = os.path.join(output_dir, \"object_collage.jpg\")\n",
    "collage(thumbnail_dir, object_collage_path)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "housekeep",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.17"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
